import numpy as np
import sys, os
import importlib
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from utils.utils import get_ref_src_path
from backends.backend_registry import BACKEND_REGISTRY

class Kernel:
    def __init__(self, op, category, language, ref_src_path=None, index=0):
        self.index = index
        self.op = op
        self.category = category
        self.language = language
        self.ref_src_path = ref_src_path
        ref_src_path = get_ref_src_path(op, category) if ref_src_path is None else ref_src_path
        with open(ref_src_path, 'r') as f:
            ref_src = f.read()
        self.ref_src = ref_src
        self.backend = self._init_backend(language)
        self.hardware = self.backend.get_hardware_name()
        self.context = {}

    def _init_backend(self, language):
        if language not in BACKEND_REGISTRY:
            try:
                importlib.import_module(f"backends.{language}_backend")
            except ImportError as e:
                raise ValueError(f"Unsupported language/platform: {language} (module not found)") from e
        backend = BACKEND_REGISTRY.get(language)
        if backend is None:
            raise ValueError(f"Unsupported language/platform: {language}")
        return backend

    def __str__(self):
        return f"{self.index} - {self.op} - {self.category}"

    def __repr__(self):
        return self.__str__()

    def set_context(self, info):
        self.context = {}
        compile(self.ref_src, "<string>", "exec")
        exec(self.ref_src, self.context)
        os.environ["ASCEND_CUSTOM_OPP_PATH"] = info['custom_opp_path']
        os.environ["LD_LIBRARY_PATH"] = info['lib_path']
        compile(info['model_src'], "<string>", "exec")
        exec(info['model_src'], self.context)

    def compile(self, code):
        """
        Compile the generated code using the backend
        Returns: (compiled: bool, compile_info: str)
        """
        if self.backend is None:
            raise ValueError("Backend not set. Please set backend before compiling.")
        
        compiled, compile_info, info = self.backend.compile(code, self.op, self.ref_src)
        result = {'index': self.index, 'op': self.op, 'category': self.category, 'language': self.language, 'compiled': compiled, 'correctness': None, 'performance': None, 'hardware':self.hardware}
        result['compiled'] = compiled
        result['compile_info'] = compile_info
        result['info'] = info
        return result

    def cleanup(self, context=None, info=None):
        """
        Clean up backend resources
        """
        if self.backend is not None:
            self.backend.cleanup(context=context, info=info)

    def execute(self, info=None):
        """
        Execute correctness and performance tests
        Returns: dict with correctness and performance results
        """
        if info:
            # 强制清理NPU缓存和同步，确保前一个kernel的资源完全释放
            import torch_npu
            torch_npu.npu.empty_cache()
            torch_npu.npu.synchronize()
            
            # 清理环境变量
            os.environ.pop("ASCEND_CUSTOM_OPP_PATH", None)
            os.environ.pop("LD_LIBRARY_PATH", None)
            
            # 设置新的环境变量
            os.environ["ASCEND_CUSTOM_OPP_PATH"] = info['custom_opp_path']
            os.environ["LD_LIBRARY_PATH"] = info['lib_path']
            
            # 重新导入custom_ops模块，确保使用新的操作符
            try:
                import importlib
                op_identifier = info.get('op_identifier')
                if op_identifier and op_identifier in sys.modules:
                    importlib.reload(sys.modules[op_identifier])
                elif op_identifier:
                    # 如果模块不存在，尝试导入
                    try:
                        importlib.import_module(op_identifier)
                    except ImportError:
                        pass
            except Exception as e:
                print(f"[WARNING] Failed to reload module: {e}")
            
            # 再次同步确保环境变量生效
            torch_npu.npu.synchronize()
        if self.backend is None:
            raise ValueError("Backend not set. Please set backend before executing.")
        
        result = {}
        
        # 如果context为空，使用backend的context
        context_to_use = self.context if self.context else self.backend.context
        
        # Execute correctness test
        correctness, correctness_info = self.backend.correctness_execution(self.ref_src, context_to_use)
        result['correctness'] = correctness
        if not correctness:
            result['correctness_info'] = correctness_info
            return result
        
        # Execute performance test
        elapsed_times = self.backend.time_execution(context=context_to_use)
        result['performance'] = {
            "mean": float(f"{np.mean(elapsed_times):.3g}"),
            "std": float(f"{np.std(elapsed_times):.3g}"),
            "min": float(f"{np.min(elapsed_times):.3g}"),
            "max": float(f"{np.max(elapsed_times):.3g}"),
            "num_trials": len(elapsed_times),
        }
        
        return result

if __name__ == "__main__":
    kernel = Kernel("mse_loss", "loss", "ascendc")
    code_path = "/home/ma-user/work/MultiKernelBench/prompts/ascendc_new_model_mse_loss.py"
    with open(code_path, 'r') as f:
        code = f.read()
    print(kernel)
    result = kernel.compile(code)
    if result['compiled']:
        kernel.set_context(result['info'])
        result.update(kernel.execute())
        print(result)
    kernel.cleanup(context=kernel.context, info=result['info'])